//
// Created by lwj12 on 24-8-31.
//

#ifndef RBTREE_RBTREE_H
#define RBTREE_RBTREE_H
#include <iostream>
using namespace std;

namespace lwj
{
    enum Color
    {
        BLACK,
        RED
    };
    template<class T>
    struct RBTreeNode
    {
        RBTreeNode* _left = nullptr;
        RBTreeNode* _right = nullptr;
        RBTreeNode* _father = nullptr;
        Color _col = RED;
        T _date;

        RBTreeNode(const T& date)
            : _date(date)
        {}
    };

    template<class T, class Ptr, class Ref>
    struct RBTIterator
    {
        typedef RBTreeNode<T> Node;
        typedef RBTIterator Self;

        Node* _node;
        Node* _root;
        RBTIterator(Node* node , Node* root)
                : _node(node),
                  _root(root)
        {}

        Ref operator*() {return _node->_date;}
        Ptr operator->() {return &_node->_date;}

        bool operator==(const Self& it) {return _node == it._node;}
        bool operator!=(const Self& it) {return _node != it._node;}

        Self& operator++()
        {
            if(_node->_right)
            {
                Node* leftMost = _node->_right;
                while(leftMost->_left) {leftMost = leftMost->_left;}
                _node = leftMost;
            }
            else if(_node->_father)
            {
                Node* father = _node->_father, *cur = _node;
                _node = nullptr;
                while(father)
                {
                    if(father->_left == cur)
                    {
                        _node = father;
                        break;
                    }
                    cur = father;
                    father = cur->_father;
                }
            }
            else _node = nullptr;
            return *this;
        }

        Self operator++(int)
        {
            Node tmp = new Node(_node, _root);
            if(_node->_right)
            {
                Node* leftMost = _node->_right;
                while(leftMost->_left) {leftMost = leftMost->_left;}
                _node = leftMost;
            }
            else if(_node->_father)
            {
                Node* father = _node->_father, *cur = _node;
                _node = nullptr;
                while(father)
                {
                    if(father->_left == cur)
                    {
                        _node = father;
                        break;
                    }
                    cur = father;
                    father = cur->_father;
                }
            }
            else _node = nullptr;
            return tmp;
        }

        Self& operator--()
        {
            if(_node->_left)
            {
                Node* rightMost = _node->_left;
                while(rightMost->_right) {rightMost = rightMost->_right;}
                _node = rightMost;
            }
            if(_node->_father)
            {
                Node* father = _node->_father, *cur = _node;
                while(father)
                {
                    if(father->_right == cur)
                    {
                        _node = father;
                        break;
                    }
                    cur = father;
                    father = cur->_father;
                }
            }
            return *this;
        }
        Self operator--(int)
        {
            Node tmp = new Node(_node, _root);
            if(_node->_left)
            {
                Node* rightMost = _node->_left;
                while(rightMost->_right) {rightMost = rightMost->_right;}
                _node = rightMost;
            }
            if(_node->_father)
            {
                Node* father = _node->_father, *cur = _node;
                while(father)
                {
                    if(father->_right == cur)
                    {
                        _node = father;
                        break;
                    }
                    cur = father;
                    father = cur->_father;
                }
            }
            return tmp;
        }
    };

    template<class K, class T, class KeyOfT>
    class RBTree
    {
    public:
        typedef RBTreeNode<T> Node;
        typedef RBTIterator<T, T*, T&> iterator;
        typedef RBTIterator<T, const T*, const T&> const_iterator;

        iterator begin()
        {
            if(_root == nullptr) return iterator(nullptr, nullptr);
            Node* cur = _root;
            while(cur->_left) cur = cur->_left;
            return iterator(cur, _root);
        }

        const_iterator end() const {return iterator(nullptr, _root);}

        const_iterator begin() const
        {
            if(_root == nullptr) return iterator(nullptr, nullptr);
            Node* cur = _root;
            while(cur->_left) cur = cur->_left;
            return iterator(cur, _root);
        }

        iterator end(){return iterator(nullptr, _root);}

        ~RBTree() {Destroy(_root); _root = nullptr;}

        void InOrder() {_InOrder(_root);cout << endl;}

        void RotateL(Node* father);
        void RotateR(Node* father);
        void RotateLR(Node* father);
        void RotateRL(Node* father);

        pair<iterator, bool> insert(const T& date);

        pair<iterator, bool> find(const T& date);

        bool check()
        {
            if(_root&&_root->_col==RED)     return false;
            Node* root = _root;
            int i = 0;
            while(root)
            {
                if(root->_col == BLACK) i++;
                root = root->_left;
            }
           return  _check(_root, i, 0);
        }

    private:
        void _InOrder(Node* root)
        {
            KeyOfT kot;
            if(root == nullptr) return;
            _InOrder(root->_left);
            cout << kot(root->_date) << " ";
            _InOrder(root->_right);
        }

        void Destroy(Node* root)
        {
            if(root == nullptr) return;
            Destroy(root->_left);
            Destroy(root->_right);
            delete root;
        }

        bool _check(Node* root, int i, int j)
        {
            if(root == nullptr)
            {
                if(i == j) return true;
                else return false;
            }

            if(root->_col == BLACK) j++;
            else
            {
                if(root->_father->_col == RED) return false;
            }
           return  _check(root->_left, i, j ) && _check(root->_right, i, j);
        }



        Node* _root = nullptr;
    };

    template<class K, class T, class KeyOfT>
    pair<typename RBTree<K, T, KeyOfT>::iterator, bool> RBTree<K, T, KeyOfT>::insert(const T &date)
    {
        Node* newNode = new Node(date);
        if(_root == nullptr) {newNode->_col = BLACK; _root = newNode; return {iterator(_root, _root), true};}
        KeyOfT kot;
        K key = kot(date);
        Node* cur = _root, *father;
        while (cur)
        {
            if(key > kot(cur->_date))
            {
                father = cur;
                cur = cur->_right;
            }
            else if(key < kot(cur->_date))
            {
                father = cur;
                cur = cur->_left;
            }
            else
            {
                delete newNode;
                return {iterator(cur, _root), false};
            }
        }

        if(kot(father->_date) > key) father->_left = newNode;
        else father->_right = newNode;
        newNode->_father = father;
        cur = newNode;

        Node* grandpa, *uncle;
        while (father && father->_col == RED)
        {
            grandpa = father->_father;
            if(grandpa->_left == father)
            {
                uncle = grandpa->_right;
                if(uncle && uncle->_col == RED)
                {
                    grandpa->_col = RED;
                    father->_col = uncle->_col = BLACK;
                    if(grandpa->_father)
                    {
                        father = grandpa->_father;
                        cur = grandpa;
                    }
                    else break;
                }
                else
                {
                       if(father->_left == cur)
                       {
                           RotateR(grandpa);
                           father->_col = BLACK;
                           grandpa->_col = RED;
                       }
                       else
                       {
                           RotateLR(grandpa);
                           cur->_col = BLACK;
                           grandpa->_col = RED;
                       }
                       break;
                }
            }
            else
            {
                uncle = grandpa->_left;
                if(uncle && uncle->_col == RED)
                {
                    grandpa->_col = RED;
                    father->_col = uncle->_col = BLACK;
                    if(grandpa->_father)
                    {
                        father = grandpa->_father;
                        cur = grandpa;
                    }
                    else break;
                }
                else
                {
                    if(father->_right == cur)
                    {
                        RotateL(grandpa);
                        father->_col = BLACK;
                        grandpa->_col = RED;
                    }
                    else
                    {
                        RotateRL(grandpa);
                        cur->_col = BLACK;
                        grandpa->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return {iterator(newNode, _root) ,true};
    }

    template<class K, class T, class KeyOfT>
    pair<typename RBTree<K, T, KeyOfT>::iterator, bool> RBTree<K, T, KeyOfT>::find(const T& date)
    {
        Node* cur = _root;
        KeyOfT kot;
        K key = kot(date);
        while (cur)
        {
            if(kot(cur->_date) > key) cur = cur->_left;
            else if(kot(cur->_date) < key) cur = cur->_right;
            else return {iterator(cur, _root), true};
        }
        return {iterator{nullptr, _root}, false};
    }

    template<class K, class T, class KeyOfT>
    void RBTree<K, T, KeyOfT>::RotateL(Node* father)
    {
        Node* subR = father->_right;
        Node* subRL = subR->_left;

        Node* fatherFather = father->_father;

        father->_father = subR;
        if(subRL) subRL->_father = father;
        father->_right = subRL;
        subR->_left = father;
        if(fatherFather)
        {
            subR->_father = fatherFather;
            if(fatherFather->_right == father) fatherFather->_right = subR;
            else fatherFather->_left = subR;
        }
        else
        {
            _root = subR;
            subR->_father = nullptr;
        }
    }

    template<class K, class T, class KeyOfT>
    void RBTree<K, T, KeyOfT>::RotateR(Node* father)
    {
        Node* subL = father->_left;
        Node* subLR = subL->_right;

        Node* fatherFather = father->_father;

        father->_father = subL;
        if(subLR) subLR->_father = father;
        father->_left = subLR;
        subL->_right = father;
        if(fatherFather)
        {
            subL->_father = fatherFather;
            if(fatherFather->_right == father) fatherFather->_right = subL;
            else fatherFather->_left = subL;
        }
        else
        {
            _root = subL;
            subL->_father = nullptr;
        }
    }

    template<class K, class T, class KeyOfT>
    void RBTree<K, T, KeyOfT>::RotateLR(Node* father)
    {
        Node* subL = father->_left, *subLR = subL->_right, *fatherFather = father->_father;

        RotateL(subL);
        RotateR(father);
    }

    template<class K, class T, class KeyOfT>
    void RBTree<K, T, KeyOfT>::RotateRL(Node* father)
    {
        Node* subR = father->_right, *subRL = subR->_left, *fatherFather = father->_father;

        RotateR(subR);
        RotateL(father);
    }


}

#endif//RBTREE_RBTREE_H
